from cProfile import label
from tkinter import Y
import pandas as pd
import numpy as np
from matplotlib import pyplot as plt

class DataDistributionGraph():
    '''
        The type of dataset must be pandas.DataFrame
    '''
    def __init__(self,dataset:pd.DataFrame,graph_type='line',label_output = False):
        self.dataset = dataset
        self.graph_type = graph_type
        self.label_output = label_output
    def draw(self):
        '''
            The output of this function include: 
                (1) Distribution of length of each sentence;
                (2) Distribution of label of each sample in dataset when label_output = True;
        '''
        len_dict = dict()
        len_dict['length'] = []
        len_dict['count'] = []
        
        if self.dataset.shape[1] == 2:
            self.dataset.columns = ['x','y']
        if self.dataset.shape[1] == 1:
            self.dataset.columns = ['x']
        for index,row in self.dataset.iterrows():
            key = len(row['x'])
            len_dict['length'].append(key)
            len_dict['count'].append(1)
        
        df = pd.DataFrame.from_dict(len_dict)
        df.columns = ['length','count']
        if self.graph_type == 'line':
            self.draw_len_line(df)
        if self.graph_type == 'bar':
            self.draw_len_bar(df)
        if self.graph_type == 'all':
            self.draw_len_line(df)
            self.draw_len_bar(df)
        if self.label_output == True:
            df = self.dataset.iloc[:,1].value_counts().plot.pie()
            plt.show()
    def draw_len_line(self,df):
        df = df.groupby('length').agg({'count':'sum'})
        df.plot()
        plt.show()
    def draw_len_bar(self,df):
        max_val = df['length'].max()
        min_val = df['length'].min()
        bins = list(i for i in range(min_val,max_val,(max_val - min_val) // 8))
        df["length"] = pd.cut(x=df["length"],bins=bins)
        df = df.dropna(axis=0)
        df = df.groupby('length').agg({'count':'sum'})
        print(df)
        print(type(df))
        df.plot(kind = 'bar')
        plt.show()

